import pymongo
from bson.objectid import ObjectId
from common.com_conf import Conf
from common.com_log import Log

class Mongo:
    __mongo_uri = f'mongodb://{Conf.yaml["leioie"]["dld"]["mongodb"]["user"]}:' \
                  f'{Conf.yaml["leioie"]["dld"]["mongodb"]["pwd"]}@' \
                  f'{Conf.yaml["leioie"]["dld"]["mongodb"]["uri"]}:' \
                  f'{Conf.yaml["leioie"]["dld"]["mongodb"]["port"]}/?' \
                  f'authSource={Conf.yaml["leioie"]["dld"]["mongodb"]["auth"]}'

    @classmethod
    def client(cls):
        """
        连接MongoDB
        :return:
        """
        return pymongo.MongoClient(cls.__mongo_uri)

    @classmethod
    def create_db(cls, db: str):
        """
        创建数据库
        :param db: 数据库名称
        :return:
        """
        cli = cls.client()
        try:
            cli[db]
            return True
        except:
            Log.error(f"mongo create_db error -> db: {db}")
            raise
        finally:
            cli.close()

    @classmethod
    def drop_db(cls, db: str):
        """
        删除数据库
        :param db: 数据库名称
        :return:
        """
        cli = cls.client()
        try:
            cli.drop_database(name_or_database=db)
            Log.info(f"mongo drop_db success -> db: {db}")
            return True
        except:
            Log.error(f"mongo drop_db error -> db: {db}")
            raise
        finally:
            cli.close()

    @classmethod
    def rename_db(cls, old_db: str, new_db: str):
        """
        修改数据库名称
        无法用
        :param old_db: 旧数据库名称
        :param new_db: 新数据库名称
        :return:
        """
        cli = cls.client()
        try:
            cli.get_database(old_db).copy_database(name_or_database=new_db)
            cli.copy_database(name_or_database=old_db)
            Log.info(f"mongo rename_db success -> db: {old_db} -> db: {new_db}")
            return True
        except:
            Log.error(f"mongo rename_db error -> db: {old_db} -> db: {new_db} ->")
            raise
        finally:
            cli.close()

    @classmethod
    def create_coll(cls, db: str, coll: str):
        """
        创建数据库表
        :param db: 数据库名称
        :param coll: 数据库集合名称
        :return:
        """
        cli = cls.client()
        try:
            cli.get_database(db)[coll]
            Log.info(f"mongo create_coll success -> db: {db} coll: {coll}")
            return True
        except:
            Log.error(f"mongo create_coll error -> db: {db} coll: {coll}")
            raise
        finally:
            cli.close()

    @classmethod
    def drop_coll(cls, db: str, coll: str):
        """
        删除数据库表
        :param db: 数据库名称
        :param coll: 数据库集合名称
        :return:
        """
        cli = cls.client()
        try:
            cli.get_database(db).drop_collection(coll)
            Log.info(f"mongo drop_coll success -> db: {db} coll: {coll}")
            return True
        except:
            Log.error(f"mongo drop_coll error -> db: {db} coll: {coll}")
            raise
        finally:
            cli.close()

    @classmethod
    def rename_coll(cls, db: str, old_coll: str, new_coll: str):
        """
        修改数据库集合名称
        :param db: 数据库名称
        :param old_coll: 数据库集合旧名称
        :param new_coll: 数据库集合新名称
        :return:
        """
        cli = cls.client()
        try:
            cli.get_database(db).get_collection(old_coll).rename(new_name=new_coll)
            Log.info(f"mongo rename_coll success -> db: {db} coll: {old_coll} -> {new_coll}")
            return True
        except:
            Log.error(f"mongo rename_coll error -> db: {db} coll: {old_coll} -> {new_coll}")
            raise
        finally:
            cli.close()

    @classmethod
    def rename_fields(cls, db: str, coll: str, filter_dict: dict = {}, rename_dict: dict = {}):
        """
        修改指定数据库的指定集合的字段键名
        :param db:
        :param coll:
        :param filter_dict:
        :param rename_dict:
        :return:
        """
        cli = cls.client()
        coll = cli.get_database(db).get_collection(coll)
        try:
            coll.update_many(filter_dict, {"$rename": rename_dict})
            Log.info(f"mongo rename_fields success -> db: {db} coll: {coll} filter: {filter_dict} update: {rename_dict}")
        except:
            Log.error(f"mongo rename_fields error -> db: {db} coll: {coll} filter: {filter_dict} update: {rename_dict}")
            raise
        finally:
            cli.close()

    @classmethod
    def find_one(cls, db, coll, cond: dict = None):
        """
        查询数据库、集合的指定文档
        :param db: 数据库名称
        :param coll: 集合
        :param cond: 条件
        :return:
        """
        cli = cls.client()
        coll = cli.get_database(db).get_collection(coll)
        try:
            if cond is None:
                cond = dict()
            Log.info(f"mongo find_one success -> db: {db} coll: {coll} cond: {cond}")
            return coll.find_one(cond)
        except:
            Log.error(f"mongo find_one error -> db: {db} coll: {coll} cond: {cond}")
            raise
        finally:
            cli.close()

    @classmethod
    def find(cls, db, coll, cond=None, sort=None, skip=0, limit=0, fields=None):
        """
        查询
        :return:
        """
        if sort is not None:
            if type([]) != type(sort):
                raise Exception("mongo find error -> sort must be array(tuple) like [('_id', -1)]")
        cli = cls.client()
        coll = cli.get_database(db).get_collection(coll)
        try:
            if cond is None:
                cond = dict()
            if 0 == skip:
                if 0 == limit:
                    if sort is None:
                        return list(coll.find(cond, fields))
                    else:
                        return list(coll.find(cond, fields).sort(sort))
                else:
                    if sort is None:
                        return list(coll.find(cond, fields).limit(limit))
                    else:
                        return list(coll.find(cond, fields).sort(sort).limit(limit))
            else:
                if 0 == limit:
                    if sort is None:
                        return list(coll.find(cond, fields).skip(skip))
                    else:
                        return list(coll.find(cond, fields).sort(sort).skip(skip))
                else:
                    if sort is None:
                        return list(coll.find(cond, fields).skip(skip).limit(limit))
                    else:
                        return list(coll.find(cond, fields).sort(sort).skip(skip).limit(limit))
        except:
            Log.error(
                f"mongo find error -> db: {db} coll: {coll} cond:{cond} sort:{sort} skip:{skip} limit:{limit} "
                f"fields:{fields} ")
            raise
        finally:
            cli.close()

    @classmethod
    def count(cls, db, coll, cond=None):
        """
        统计集合中的文档数量
        """
        cli = cls.client()
        coll = cli.get_database(db).get_collection(coll)
        try:
            if cond is None:
                cond = dict()
            return list(coll.find(cond)).count()
        except:
            Log.error(f"mongo count error -> db: {db} coll: {coll} cond: {cond}")
            return None
        finally:
            cli.close()

    @classmethod
    def update_one(cls, db, coll, filter_dict, update_dict, upsert=True):
        """
        更新指定文档
        """
        cli = cls.client()
        coll = cli.get_database(db).get_collection(coll)
        try:
            coll.update_one(filter=filter_dict, update={"$set": update_dict}, upsert=upsert)
            Log.info(f"mongo update_one success -> db: {db} coll: {coll} filter: {filter_dict} update: {update_dict}")
        except:
            Log.error(f"mongo update_one error -> db: {db} coll: {coll} filter: {filter_dict} update: {update_dict}")
            raise
        finally:
            cli.close()

    @classmethod
    def update_many(cls, db, coll, filter_dict, update_dict, upsert=True):
        """
        更新文档
        :param db:
        :param coll:
        :param filter_dict:
        :param update_dict:
        :param upsert:
        :return:
        """
        cli = cls.client()
        coll = cli.get_database(db).get_collection(coll)
        try:
            coll.update_many(filter=filter_dict, update={"$set": update_dict}, upsert=upsert)
            Log.info(f"mongo update_many success -> db: {db} coll: {coll} filter: {filter_dict} update: {update_dict}")
        except:
            Log.error(f"mongo update_many error -> db: {db} coll: {coll} filter: {filter_dict} update: {update_dict}")
            raise
        finally:
            cli.close()

    @classmethod
    def rm_insert_many(cls, db, coll, dict_arr: dict):
        """
        先删除原表，再插入数据
        :param db:
        :param coll:
        :param dict_arr:
        :return:
        """
        cli = cls.client()
        coll = cli.get_database(db).get_collection(coll)
        try:
            coll.drop()
            coll.insert_many(dict_arr, ordered=False)
            Log.info(f"mongo rm_insert_many success -> db: {db} coll: {coll}")
        except:
            Log.error(f"mongo rm_insert_many error -> db: {db} coll: {coll}")
            raise
        finally:
            cli.close()

    @classmethod
    def insert_many(cls, db, coll, dict_arr):
        """
        插入数据
        :param db:
        :param coll:
        :param dict_arr:
        :return:
        """
        cli = cls.client()
        coll = cli.get_database(db).get_collection(coll)
        try:
            coll.insert_many(dict_arr, ordered=False)
            Log.info(f"mongo insert_many success -> db: {db} coll: {coll}")
        except:
            Log.error(f"mongo insert_many error -> db: {db} coll: {coll}")
            raise
        finally:
            cli.close()

    @classmethod
    def delete_many(cls, db, coll, filter_dict):
        """
        删除指定文档
        :param db   数据库名称
        :param coll 集合名称
        :param filter_dict  查询条件
        """
        cli = cls.client()
        coll = cli.get_database(db).get_collection(coll)
        try:
            coll.delete_many(filter=filter_dict)
            Log.info(f"mongo delete_many success -> db: {db} coll: {coll} filter: {filter_dict}")
        except:
            Log.error(f"mongo delete_many error -> db: {db} coll: {coll} filter: {filter_dict}")
            raise
        finally:
            cli.close()


if __name__ == '__main__':
    # print(Mongo.find(db='stock_1min', coll="000001.SZ", cond={'trade_date': {'&glt': '20230101'}}))
    """
    mongodump -h localhost -d basic -o D:/MongoDBData --port 27017 -u admin -p 123456 --authenticationDatabase admin
    mongodump -h 127.0.0.1 --forceTableScan -d basic -o D:/MongoDBData --port 27017 -u admin -p 123456 --authenticationDatabase admin
    mongodump mongodb://admin:123456@127.0.0.1:27017/?authSource=admin  -d basic -o D:/MongoDBData
    """
    # print(Mongo.find(db='basic', coll="hello", cond={'name': 'Tom'}))
    # # Mongo.update_one(db='basic', coll='hello', filter_dict={'name': 'Jack'}, update_dict={'gender': 'Man'})
    # # Mongo.update_one(db='basic', coll='hello', filter_dict={'name': 'Jack'}, update_dict={'gender': 'Man'})
    # Mongo.update_many(db='basic', coll='hello', filter_dict={}, update_dict={'hobby': 'Singer'})

    # stock_basic = Mongo.find(db='basic', coll='stock_basic')
    # for item in stock_basic:
    #     if '.' in item['ts_code']:
    #         print(item['ts_code'])
    #         Mongo.delete_many(db='stock_1min', coll=item['ts_code'],
    #                           filter_dict={'trade_date': {'$gte': '20230721', '$lte': '20230721'}})

    # stock_basic = Mongo.find(db='basic', coll='stock_basic')
    # print(stock_basic)
    # index = [i for i in range(len(stock_basic)) if stock_basic[i]['ts_code'] == '688262.SH'][0]
    # # print(index)
    # for item in stock_basic[index:]:
    #     if '.' in item['ts_code']:
    #         print(item['ts_code'])
    # #         Mongo.update_many(db='stock_1min', coll=item['ts_code'], filter_dict={},
    # #                           update_dict={'name': item['name']})

    # print(Mongo.create_db('school'))
    # print(Mongo.create_coll('school', 'student'))
    # # # print(Mongo.client().get_database('basic').create_collection(name='trade_calendar'))
    # Mongo.insert_many(db='school', coll='student',
    #                   arr=[{'_id': '20230515', 'open': 16}, {'_id': '20230516', 'open': 18}])
    # print(Mongo.find(db="basic", coll="stock_basic", cond={'ts_code': '600519.SH'}))

    # print(Mongo.find(db="basic", coll="stock_basic", cond={}))
    for item in Mongo.find(db="basic", coll="stock_basic", cond={})[:]:
        if '.BJ' not in item['ts_code']:
            print(item)
            Mongo.update_many(db='stock_5min_qfq', coll=f"{item['ts_code']}", filter_dict={},
                              update_dict={'ts_code': item['ts_code']})
            # Mongo.rename_coll(db="stock_daily_qfq", old_coll=f"qfq_{item['ts_code']}", new_coll=f"{item['ts_code']}")
        # print('分钟数据：', Mongo.find(db='stock_1min', coll=item['ts_code'], cond={'_id': {'$gte': '2023-05-16', '$lte': '2023-05-17'}}))
        # print('分钟数据：', Mongo.delete_many(db='stock_1min', coll=item['ts_code'], filter_dict={'_id': {'$gte': '2023-12-01', '$lte': '2023-12-02'}}))
    # Mongo.delete_many(db='basic', coll="trade_calendar",
    #                   filter_dict={'cal_date': {'$gte': '20230101', '$lte': '20230619'}})

    # # print(Mongo.find(db='basic', coll='trade_calendar', cond={'trade_date': '20230514'}))
    # print(Mongo.find(db='basic', coll='trade_calendar', cond={'_id': ObjectId('646192132f60bbf4b5360e37')}))
    # Mongo.delete_many(db='basic', coll='trade_calendar', filter_dict={'_id': ObjectId('646192132f60bbf4b5360e37')})
    # print(Mongo.find(db='basic', coll='trade_calendar'))
